import os
import hashlib
import shutil
import requests
from tqdm import tqdm
import time
import json

md5Path = r"/mnt/mydisk1/chenyongwei/malware/PEDatas/md5Datas"
tmpPath = r"/mnt/mydisk1/chenyongwei/malware/tmpDatas"
PEPath = r"/mnt/mydisk1/chenyongwei/malware/PEDatas"
dataPath = r"/mnt/mydisk1/chenyongwei/malware/dailyDatas"
basicInfPath = r"/mnt/mydisk1/chenyongwei/malware/PEDatas/basicInfDatas"
famLablePath = r"/home/cyw/projects/function_sim_project/malware_bazaar_dataset/famLable.txt"


def selectPEFile():
    """
        提取将下载样本的PE文件，并将其重命名为md5值
    """
    files = os.listdir(dataPath)
    for file in tqdm(files):
        tmp = file.split(".")
        if (len(tmp) == 2 and tmp[1] == "exe"):
            with open(dataPath+"//"+file, 'rb') as fp:
                data = fp.read()
            fileMd5 = hashlib.md5(data).hexdigest()
            shutil.copy(dataPath+"//"+file, PEPath+"//"+file)
            shutil.copy(dataPath+"//"+file, md5Path+"//"+file)
            os.rename(md5Path+"//"+file, md5Path+"//"+fileMd5)
            with open("/home/cyw/projects/function_sim_project/malware_bazaar_dataset/PEtransLog.exe", 'a') as fp:
                fp.write(file+"\t--->\t"+fileMd5+"\n")


def querySampleBasicInfByMd5(md5):
    """
        请求mb，获得样本的基础信息
    """
    #  wget --header "Auth-Key: YOUR-AUTH-KEY-HERE" --post-data "query=get_info&hash=7de2c1bf58bce09eecc70476747d88a26163c3d6bb1d85235c24a558d1f16754" https://mb-api.abuse.ch/api/v1/
    url = "https://mb-api.abuse.ch/api/v1/"
    headers = {
        "Auth-Key": "1c830b4d62fa0c9b0f35229d080605c42082049270bf49e8"
    }
    data = {
        "query": "get_info",
        "hash": md5
    }
    response = requests.post(url, headers=headers, data=data)

    str = md5+" "
    if response.status_code == 200:
        with open(basicInfPath+"//"+"{}.html".format(md5), "wb") as file:
            file.write(response.content)
        str += "文件下载成功！\n"
    else:
        str += "请求失败，信息：{}\n".format(response)

    with open("/home/cyw/projects/function_sim_project/malware_bazaar_dataset/querySampleInfLog.txt", "a") as file:
        file.write(str)


def generateAllsampleInf():
    """
        利用mb,获得所有样本的信息
    """
    # 样本太多，支持中断
    allSamples = os.listdir(md5Path)
    tmps = os.listdir(basicInfPath)
    finishedSamples = {}
    for tmp in tmps:
        name = tmp.split(".")[0]
        finishedSamples[name] = True
    unfinishSamples = []
    for sample in allSamples:
        if sample not in finishedSamples:
            unfinishSamples.append(sample)
    print(len(allSamples))
    for sample in tqdm(unfinishSamples):
        time.sleep(15)
        querySampleBasicInfByMd5(sample)


def generateSampleFamLable():
    """
        生成样本名和其对应的家族标签
    """
    files = os.listdir(basicInfPath)
    for file in files:
        tmp = file.split(".")
        sampleName = tmp[0]
        with open(basicInfPath+"//"+sampleName+".html", 'r') as sampleBasicInf:
            data = sampleBasicInf.read()
        htmlData = json.loads(data)
        try:
            sampleLable = str(htmlData["data"][0]["signature"])
            if sampleLable == "None":
                raise Exception('不存在家族标签')
            with open(famLablePath, 'a') as saveFile:
                saveFile.write(str(sampleName)+"\t" + sampleLable + "\n")
        except Exception as e:
            print("error" + sampleName + "\t" + str(e))


def getSample2Fam():
    """
        获得样本名和其对应的家族标签
    """
    sample2Fam = {}
    with open(famLablePath, 'r') as file:
        for line in file.readlines():
            tmp = line.strip().split("\t")
            sample = tmp[0]
            lable = tmp[1]
            sample2Fam[sample] = lable
    return sample2Fam


def generatePEWithLable():
    """
        生成转移到window中，提供给ida pro处理的PE软件
    """
    sample2Fam = getSample2Fam()
    cnt = 10
    for sample in tqdm(sample2Fam):
        shutil.copy(md5Path+"//"+sample, tmpPath+"//"+sample)
        cnt -= 1
        if cnt == 0:
            break


if __name__ == "__main__":
    # selectPEFile()
    generateAllsampleInf()
    # generateSampleFamLable()
    # generatePEWithLable()
